import pickle

import akshare
import numpy
import tqdm


def readStockConfig(config_path):
    # 股票列表
    stock_list = []
    if config_path is None:
        config_path = './DATA/invest_opt/stock.config'

    with open(config_path) as f:
        while (True):
            try:
                tmp = f.readline()
                stock_list.append(int(tmp))
            except:
                break
    return stock_list


def getStockData(start_date='20200101', end_date='20210831', config_path=None, cache_path=None):
    stock_codes_list = readStockConfig(config_path)

    if cache_path is None:
        cache_path = './DATA/invest_opt/invest.cache'

    try:
        with open(cache_path, 'rb') as f:
            e = pickle.load(f)
    except:
        # 开盘价
        e = []
        # 获取数据
        for i in tqdm.tqdm(stock_codes_list):
            # 获取指定时间的交易数据
            a = akshare.stock_zh_a_hist(symbol=str(i), period="daily", start_date=start_date,
                                        end_date=end_date, adjust="")
            e.append(a['收盘'])
        # 转numpy
        maxlen = numpy.min([i.shape[-1] for i in e])
        e = numpy.vstack([i[-maxlen:] for i in e])

        with open(cache_path, 'wb') as f:
            pickle.dump(e, f)
    finally:
        t = numpy.log(e[:, 1:] / e[:, :-1] * (1 - .00132))
        return t
